package com.yubo.yupicturebackend.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yubo.yupicturebackend.common.BaseResponse;
import com.yubo.yupicturebackend.common.ResultUtils;
import lombok.Data;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

/**
 * ai聊天接口
 * @author Seryer
 */
@RestController
@RequestMapping("/ai")
public class AiChatController {

    private final WebClient webClient;

    public AiChatController(WebClient.Builder webClientBuilder) {
        this.webClient = webClientBuilder.baseUrl("http://localhost:11434")
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE).build();
    }

    @PostMapping(value = "/stream-chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<BaseResponse<String>> streamChat(@RequestBody OllamaRequest request) {
        return webClient.post()
                .uri("/api/generate")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(request)
                .retrieve()
                // 接收 Ollama 的流式响应
                .bodyToFlux(String.class)
                // 提取响应字段
                .map(this::extractResponse);
    }

    /**
     *  解析 Ollama 的 JSON 块，提取 "response" 字段
     * @param jsonChunk
     * @return
     */
    private BaseResponse<String> extractResponse(String jsonChunk) {
//        try {
//            ObjectMapper mapper = new ObjectMapper();
//            JsonNode node = mapper.readTree(jsonChunk);
//            return node.has("response") ? node.get("response").asText() : "";
//        } catch (Exception e) {
//            return "[解析错误]";
//        }
        return ResultUtils.success(jsonChunk);
    }


    /**
     * 请求体
     */
    @Data
    public static class OllamaRequest {
        private String model;
        private String prompt;
        private boolean stream = true;

        // getters/setters
    }
}